from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Random import get_random_bytes
import base64


# 生成 RSA 密钥对
def generate_keys():
    key = RSA.generate(2048)
    private_key = key.export_key()
    public_key = key.publickey().export_key()

    # 将私钥保存到文件
    with open('mqtt_private.pem', 'wb') as f:
        f.write(private_key)

    # 将公钥保存到文件
    with open('mqtt_public.pem', 'wb') as f:
        f.write(public_key)

    return private_key, public_key

generate_keys()
# 使用公钥加密
# def encrypt_message(message, public_key):
#     key = RSA.import_key(public_key)
#     cipher = PKCS1_v1_5.new(key)
#     ciphertext = cipher.encrypt(message.encode())
#     return base64.b64encode(ciphertext)
#
#
# # 使用私钥解密
# def decrypt_message(ciphertext, private_key):
#     key = RSA.import_key(private_key)
#     cipher = PKCS1_v1_5.new(key)
#     message = cipher.decrypt(base64.b64decode(ciphertext), None)
#     return message.decode()
#
#
# # 生成密钥对并保存到文件
# # private_key, public_key = generate_keys()
# private_key = open('private.pem', 'rb').read()
# public_key = open('public.pem', 'rb').read()
# # 要加密的消息
# message = "Hello, this is a secret message."
#
# # 使用公钥加密消息
# encrypted_message = encrypt_message(message, public_key)
# print("Encrypted message:", encrypted_message)
#
# # 使用私钥解密消息
# decrypted_message = decrypt_message(encrypted_message, private_key)
# print("Decrypted message:", decrypted_message)
